import torch
import matplotlib.pyplot as plt
from torch import nn  # 构建神经网络的模型框架

# 未知的规律
# -----------------散点图------------------
x = torch.linspace(-2, 2, 40)
y = 3 * x ** 2 + 1

# print(x.shape)  # 40长度
# print(y.shape)  # 40长度
x = x.reshape(40, 1)  # 由于模型训练需要输入一个特征，因此最后一个维度必须是1
y = y.reshape(40, 1)  # 40 --> batch_size 批次大小

plt.plot(x, y, 'ro')

# -----------------预测线--------------------
w_predict = torch.tensor(0.1, requires_grad=True)  # 预测的w
b_predict = torch.tensor(0.2, requires_grad=True)  # 预测的b
# -----------------构建模型--------------------
model = nn.Sequential(
    nn.Linear(1, 10),  # 画直线
    nn.Sigmoid(),  # 为整体加入激活函数，让直线弯曲
    nn.Linear(10, 1)  # 曲线进行缩放和平移
)
# 绘制预测线
y_predict = model(x)
line, = plt.plot(x.detach().numpy(), y_predict.detach().numpy(), 'b--')
# ------------------损失函数---------------------
criterion = nn.MSELoss()  # 均方差损失
# ------------------梯度下降---------------------
# model.parameters() 整个模型参与运算的w和b
# lr = learning_rate 学习率
sgd = torch.optim.SGD(model.parameters(), lr=0.05)
# ------------------训练---------------------
epochs = 1000
for epoch in range(epochs):
    # 清空所有参与的w，b的导数
    sgd.zero_grad()
    # 前向传播：使用模型预测
    y_predict = model(x)
    # 重绘整个图
    line.set_data(x.detach().numpy(), y_predict.detach().numpy())
    # 损失函数
    loss = criterion(y_predict, y)
    # 反向传播
    loss.backward()
    sgd.step()

    if epoch % 10 == 0:
        print(f"epoch {epoch + 1}/{epochs} -- loss:{loss.item():.4f}")
    # 设置显示延迟
    plt.pause(0.1)
